import pandas as pd
from django.db import transaction
import logging

logger = logging.getLogger(__name__)


def import_excel_to_db(model_class, excel_path='./data.xlsx', sheet_name=0, mapping=None):
    """
    通用的 Excel 导入数据库工具方法

    参数:
        model_class: Django 模型类
        excel_path: Excel 文件路径
        sheet_name: 工作表名称或索引（默认为第一个工作表）
        mapping: 字典，用于将 Excel 列名映射到模型字段名
                例如: {'Excel列名': 'model字段名'}

    返回:
        tuple: (成功导入的记录数, 错误信息列表)
    """
    try:
        # 读取 Excel 文件
        df = pd.read_excel(excel_path, sheet_name=sheet_name)

        # 如果提供了映射，重命名列
        if mapping:
            df = df.rename(columns=mapping)

        # 获取模型的字段名
        model_fields = [f.name for f in model_class._meta.fields if f.name != 'id']

        # 过滤出模型中存在的列
        valid_columns = [col for col in df.columns if col in model_fields]
        df = df[valid_columns]

        success_count = 0
        error_messages = []

        # 使用事务进行批量导入
        with transaction.atomic():
            # 将 DataFrame 转换为字典列表
            records = df.to_dict('records')

            # 批量创建对象
            objects_to_create = []
            for record in records:
                try:
                    # 创建模型实例
                    obj = model_class(**record)
                    objects_to_create.append(obj)
                except Exception as e:
                    error_messages.append(f"处理记录时出错: {str(e)}")

            if objects_to_create:
                # 批量创建记录
                model_class.objects.bulk_create(objects_to_create)
                success_count = len(objects_to_create)

        return success_count, error_messages

    except Exception as e:
        return 0, [f"导入过程出错: {str(e)}"]


def import_excel_with_progress(model_class, excel_path='./data.xlsx', sheet_name=0, mapping=None, batch_size=1000):
    """
    带进度显示的 Excel 导入数据库工具方法

    参数:
        model_class: Django 模型类
        excel_path: Excel 文件路径
        sheet_name: 工作表名称或索引
        mapping: 字典，用于将 Excel 列名映射到模型字段名
        batch_size: 每批处理的记录数

    返回:
        tuple: (成功导入的记录数, 错误信息列表)
    """
    try:
        # 读取 Excel 文件
        df = pd.read_excel(excel_path, sheet_name=sheet_name)
        total_rows = len(df)

        if mapping:
            df = df.rename(columns=mapping)

        model_fields = [f.name for f in model_class._meta.fields if f.name != 'id']
        valid_columns = [col for col in df.columns if col in model_fields]
        df = df[valid_columns]

        success_count = 0
        error_messages = []

        # 分批处理数据
        for start_idx in range(0, total_rows, batch_size):
            end_idx = min(start_idx + batch_size, total_rows)
            batch_df = df.iloc[start_idx:end_idx]

            with transaction.atomic():
                objects_to_create = []
                for _, row in batch_df.iterrows():
                    try:
                        obj = model_class(**row.to_dict())
                        objects_to_create.append(obj)
                    except Exception as e:
                        error_messages.append(f"处理记录时出错: {str(e)}")

                if objects_to_create:
                    model_class.objects.bulk_create(objects_to_create)
                    success_count += len(objects_to_create)

            # 打印进度
            progress = (end_idx / total_rows) * 100
            logger.info(f"导入进度: {progress:.2f}% ({end_idx}/{total_rows})")

        return success_count, error_messages

    except Exception as e:
        return 0, [f"导入过程出错: {str(e)}"]
